package cn.iocoder.yudao.module.dify.service.BatchUploadKnowledgeBase;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.file.FileNameUtil;
import cn.iocoder.yudao.module.dify.util.DocToDocxConverter;
import lombok.extern.slf4j.Slf4j;
import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import org.apache.pdfbox.io.MemoryUsageSetting;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.nio.file.Paths;

@Slf4j
@Service
public class PDFProcessService {

    @Autowired
    private ITesseract tesseract;

    @Autowired
    private BatchUploadKnowledgeBaseService batchUploadKnowledgeBaseService;

    // 生产者 - 将OCR任务放入Redis队列
    @Autowired
    @Qualifier("stringRedisTemplate")  // 或 "justAuthRedisCacheTemplate"
    private RedisTemplate<String, String> redisTemplate;

    private static final String OCR_TASK_QUEUE = "ocr:task:queue";
    private static final String ERROR_OCR_TASK_QUEUE = "error:task:queue";
    private static final String DOCX_TASK_QUEUE = "docx:task:queue";

    public void addFileToOCRQueue(File file, String datasetId) {
        // 序列化任务信息
        String taskInfo = file.getAbsolutePath() + "|" + datasetId;
        redisTemplate.opsForList().leftPush(OCR_TASK_QUEUE, taskInfo);
    }

    public void addFileToDocxQueue(File file, String datasetId) {
        // 序列化任务信息
        String taskInfo = file.getAbsolutePath() + "|" + datasetId;
        redisTemplate.opsForList().leftPush(DOCX_TASK_QUEUE, taskInfo);
    }

    // 消费者 - 从Redis队列取出任务并处理
    @Scheduled(fixedDelay = 2000) // 或使用其他方式启动监听
    public void processOCRTasks() {
        String taskInfo = redisTemplate.opsForList().rightPop(OCR_TASK_QUEUE);
        if (taskInfo != null) {
            try {
                String[] parts = taskInfo.split("\\|");
                File file = new File(parts[0]);
                String datasetId = parts[1];
                // 执行OCR处理逻辑
                batchUploadKnowledgeBaseService.processFileWithOCR(file, datasetId);
                log.info("处理文件: " + file.getAbsolutePath());
            } catch (Exception e) {
                // 处理失败时记录日志或重新入队
                log.error("OCR任务处理失败: " + taskInfo, e);
                // 放入失败队列
                redisTemplate.opsForList().leftPush(ERROR_OCR_TASK_QUEUE, taskInfo);
            }
        }
    }

    @Scheduled(fixedDelay = 3000) // 或使用其他方式启动监听
    public void docToDocxTasks() throws Exception {
        String taskInfo = redisTemplate.opsForList().rightPop(DOCX_TASK_QUEUE);
        if (taskInfo != null) {
            try {
                String[] parts = taskInfo.split("\\|");
                File file = new File(parts[0]);
                String datasetId = parts[1];
                // 执行docx处理逻辑
                File docxFile = new File(file.getAbsoluteFile() + "x");
                DocToDocxConverter.convertDocToDocx(file, docxFile);
                batchUploadKnowledgeBaseService.uploadFileToKnowledgeBase(docxFile, datasetId);
                log.info("处理doc文件成功: " + docxFile.getAbsolutePath());
            } catch (Exception e) {
                // 处理失败时记录日志或重新入队
                log.error("OCR任务处理失败: " + taskInfo, e);
                // 放入失败队列
                redisTemplate.opsForList().leftPush(ERROR_OCR_TASK_QUEUE, taskInfo);
            }
        }
    }

    /**
     * 智能处理PDF文件
     *
     * @param pdfFile PDF文件
     * @return 处理结果
     */
    public PDFProcessResult processPDF(File pdfFile) {
        PDFProcessResult result = new PDFProcessResult();

        try {
            // 尝试提取文本
            String extractedText = extractText(pdfFile);

            if (isTextValid(extractedText)) {
                result.setSuccess(true);
                result.setText(extractedText);
                result.setUsingOCR(false);
            } else {
                // 文本提取失败，使用OCR
                File ocrText = performOCR(pdfFile);
                result.setSuccess(true);
                result.setText(FileUtil.readUtf8String(ocrText));
                result.setUsingOCR(true);
            }
        } catch (Exception e) {
            result.setSuccess(false);
            result.setError(e.getMessage());
        }
        return result;
    }

    private String extractText(File pdfFile) throws IOException {
        try (PDDocument document = PDDocument.load(pdfFile, MemoryUsageSetting.setupTempFileOnly())) {
            // 检查文档是否加密
            if (document.isEncrypted()) {
                log.warn("PDF文件已加密，跳过文本提取: {}", pdfFile.getAbsolutePath());
                return "";
            }
            PDFTextStripper stripper = new PDFTextStripper();
            stripper.setStartPage(1);
            stripper.setEndPage(document.getNumberOfPages());
            return stripper.getText(document).trim();
        } catch (Exception e) {
            log.error("提取PDF文本失败: " + pdfFile.getAbsolutePath(), e);
            return ""; // 返回空字符串而不是抛出异常
        }
    }

    private boolean isTextValid(String text) {
        // 简单验证文本是否有效（可以根据需求调整）
        return text != null && text.length() > 3;
    }

    public File performOCR(File pdfFile) throws Exception {
        // 检查文件
        if (!pdfFile.exists() || !pdfFile.canRead()) {
            throw new IOException("PDF文件不存在或不可读: " + pdfFile.getAbsolutePath());
        }
        StringBuilder ocrText = new StringBuilder();
        try (PDDocument document = PDDocument.load(pdfFile)) {
            PDFRenderer renderer = new PDFRenderer(document);
            for (int i = 0; i < document.getNumberOfPages(); i++) {
                BufferedImage image = renderer.renderImageWithDPI(i, 300);
                String pageText = tesseract.doOCR(image);
                ocrText.append(pageText).append("\n\n");
            }
        }
        return saveTextToTxt(ocrText.toString(), pdfFile);
    }

    /**
     * 执行OCR识别，添加同步机制避免并发访问问题
     *
     * @param pdfFile PDF文件
     * @return 生成的txt文件
     * @throws Exception
     */
    public synchronized File performOCRSave(File pdfFile) throws Exception {
        StringBuilder ocrText = new StringBuilder();

        try (PDDocument document = PDDocument.load(pdfFile)) {
            PDFRenderer renderer = new PDFRenderer(document);

            // 为当前OCR任务创建独立的Tesseract实例
            ITesseract localTesseract = new Tesseract();
            URL tessDataURL = getClass().getClassLoader().getResource("tessdata");
            if (tessDataURL == null) {
                throw new RuntimeException("Cannot find tessdata directory in resources");
            }
            String tessDataPath = Paths.get(tessDataURL.toURI()).toString();
            localTesseract.setDatapath(tessDataPath);
            // 设置默认语言
            localTesseract.setLanguage("chi_sim+eng");

            // 设置OCR引擎模式
            localTesseract.setOcrEngineMode(1); // 使用LSTM OCR引擎

            // 设置页面分割模式
            localTesseract.setPageSegMode(6); // 假设一个统一的文本块
            for (int i = 0; i < document.getNumberOfPages(); i++) {
                BufferedImage image = renderer.renderImageWithDPI(i, 300);
                String pageText = localTesseract.doOCR(image);
                ocrText.append(pageText).append("\n\n");
            }
        }
        // 保存文本到txt文件并返回文件对象
        String txtFileName = pdfFile.getAbsolutePath().replace(".pdf", ".txt");
        File txtFile = new File(txtFileName);
        FileUtil.writeUtf8String(ocrText.toString(), txtFile);
        return txtFile;
    }

    /**
     * 判断是否需要OCR
     *
     * @param pdfFile
     * @return
     * @throws IOException
     */
    public boolean needOCR(File pdfFile) {
        if (!FileNameUtil.getSuffix(pdfFile).equals("pdf")) {
            return false;
        }
        int maxRetries = 3;
        for (int i = 0; i < maxRetries; i++) {
            try {
                String text = extractText(pdfFile);
                return !isTextValid(text);
            } catch (Exception e) {
                if (i == maxRetries - 1) {
                    log.error("多次尝试后仍无法处理PDF文件: " + pdfFile.getAbsolutePath(), e);
                    return false;
                }
            }
        }
        return false;
    }

    /**
     * 文本存入同路径文件夹同名txt文件
     *
     * @param text
     * @param pdfFile
     * @return txt文件路径
     */
    public File saveTextToTxt(String text, File pdfFile) {
        String txtFileName = pdfFile.getAbsolutePath().replace(".pdf", ".txt");
        File txtFile = new File(txtFileName);
        //text写入txt文件
        FileUtil.writeUtf8String(text, txtFile);
        return txtFile;
    }
}
